from utils import INF, is_sorted

from khronos.opt.mh.treesearch.sts import STS

class TargetSTS(STS):
    """Target-oriented stochastic tree search. Target-oriented search is an extension to the 
    regular STS algorithm. The idea is, after an improving solution is found, to search the 
    queue of open nodes for partial solutions which are similar to neighbor solutions of the 
    most recent solution, and make a dive from the most similar partial solution in order to 
    obtain the neighbor (target) at the end of the dive.
    Using this class one can easily extend the STS algorithm with local search and path 
    relinking, simply by defining the following methods:
        targets(solution)          - return a list of target neighbor solutions from the given 
                                     solution
        similarity_to_target(node) - compute the similarity of the partial solution contained in
                                     the given node to the current target solution
        build_to_target(node)      - complete the partial solution in a node in order to obtain 
                                     the current target solution"""
    __slots__ = ["target"]
    def __init__(self, *args, **kwargs):
        super(TargetSTS, self).__init__(*args, **kwargs)
        self.target = None
        
    def bootstrap(self):
        super(TargetSTS, self).bootstrap()
        self.target = None
        
    def add_solution(self, solution, z):
        super(TargetSTS, self).add_solution(solution, z)
        self.target_search(solution)
        
    def target_search(self, pivot):
        """Perform a target search from a given pivot solution (usually the most recently 
        obtained solution). Target search consists of three steps. Initially, a set of target 
        solutions is obtained from the initial solution, e.g. a set of neighbor solutions, as 
        in local search. These target solutions will guide the target search in the subsequent 
        steps. Then, for each target solution, the queue of open nodes is searched for the 
        partial solution with more resemblance to the target solution. Finally, the most 
        similar node in the queue is explored, for instance by making a full-expansion dive 
        which completes the partial solution as the target solution. The node is removed from 
        the queue afterwards if all its children were generated and analyzed, otherwise it may 
        still lead to the optimal solution, so it is unsafe to discard it.
        This kind of procedure prioritizes such search methods as local search and/or path 
        relinking because, although keeping tree search's avoidance of repeated search, it 
        explores neighboring partial solutions potentially before their queue rank would dictate. 
        Although expensive, this type of search can lead to an early tightening of the upper 
        bound, thus yielding significant speed gains later on in the search. Furthermore, target 
        search, when the list of targets represents some neighborhood of the search's pivot 
        solution, ***leads to local optimum solutions***! This guarantee cannot be given by 
        regular STS unless local search is externally executed."""
        # Exit immediately if another target search is already in progress.
        if self.target is not None:
            return
        initial_UB  = self.upper_bound
        iterations  = 0
        while pivot is not None:
            nsolutions = len(self.solutions)
            previous_UB = self.upper_bound
            iterations += 1
            # Begin by generating target solutions from the current pivot solution.
            for self.target in self.targets(pivot):
                # Continue by finding the node with maximum similarity to each target. 
                # If no node was found with positive similarity we consider the target 
                # already explored and skip it, otherwise we complete the partial solution 
                # contained in the closest node. The node is removed from the queue if 
                # build_to_target() returns True, meaning that it was fully expanded and 
                # can be safely removed.
                closest = self.closest_to_target()
                if closest is not None:
                    fully_expanded = self.build_to_target(closest.node)
                    if fully_expanded:
                        self.queue.outsort(closest)
                        assert closest not in self.queue
                        assert is_sorted(self.queue)
            self.target = None
            if len(self.solutions) > nsolutions:
                pivot = self.solutions[-1].Solution
            else:
                pivot = None
            self.log("TARGET SEARCH: %s -> %s (iteration %d)" % 
                     (previous_UB, self.upper_bound, iterations))
        self.log("TARGET SEARCH: %s -> %s (%d iterations)" % 
                 (initial_UB, self.upper_bound, iterations))
        
    def closest_to_target(self):
        """Search for the node in the queue which contains the partial solution with most 
        similarity to the current target solution."""
        closest = None
        max_similarity = (0.0, 1.0)
        for wrapper in self.queue:
            s = self.similarity_to_target(wrapper.node)
            if s > 0 and closest is not None:
                raise Exception("more than one similar node")
            similarity = (s, self.rng.random())
            if similarity > max_similarity:
                closest = wrapper
                max_similarity = similarity
                if s == INF:
                    break
        return closest
        
    # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    def targets(self, pivot): # OPTIONAL OVERRIDE
        """Return a list of solutions to be set as the next targets for the search."""
        return ()
        
    def similarity_to_target(self, node): # MANDATORY OVERRIDE
        """Compute and return a value expressing how similar an open node containing a partial 
        solution is to the current target solution. If this method returns INF (positive infinity) 
        similarity for any node, the search for a node with higher similarity to the target is 
        stopped."""
        raise NotImplementedError()
        
    def build_to_target(self, node): # MANDATORY OVERRIDE
        """Complete the given partial solution until it is equal to the current target solution. 
        This can be easily done through a proper call to the dive() method, where a subclass can
        define the dive's path and expand functions and also the abort value. This method must 
        return a boolean value indicating if the given node was fully expanded, to allow 
        target_search() to remove it from the queue in that case."""
        raise NotImplementedError()
        
